import time
import os
import pylab as plt
import sys
import datetime
import math
import matplotlib.dates as mdates

from matplotlib.pyplot import figure, show
from matplotlib.dates import WeekdayLocator#, HourLocator, DateFormatter, drange
from numpy import arange

from params import *

tracefile = os.path.join(WORKDIR, 'gowalla', 'temporal', 'data',
    'gowalla_growth_data.txt')
days = []
nodes = []
edges = []
degrees = []
for line in open(tracefile):
    tokens = line.split(' ')
    date = tokens[0]
    d,m,y = map(int,date.split('-'))
    date = datetime.datetime(y,m,d)
    n,k= map(int,(tokens[1:3]))
    deg = float(tokens[3])

    days.append(date)
    nodes.append(n)
    degrees.append(deg)
    edges.append(k)

major_loc = mdates.WeekdayLocator(byweekday=1, interval=2)
minor_loc = mdates.WeekdayLocator(byweekday=1)
fmtr = mdates.DateFormatter('%d %b')

#x = days
#y1 = nodes
#y2 = edges
#fig = plt.figure()
#plt.clf()
#ax = plt.axes(FIG_AXES2)
#plt.plot_date(x,y1,fmt='k-')
#plt.plot_date(x,y2,fmt='k--')
##plt.ylabel(title)
#plt.legend(['Nodes', 'Links'])
#plt.grid(True)
#labels = ax.get_xticklabels()
#plt.setp(labels, rotation=30, fontsize=8)
#ax.xaxis.set_major_locator(major_loc)
#ax.xaxis.set_minor_locator(minor_loc)
#ax.xaxis.set_major_formatter(fmtr)
#plt.savefig('gowalla_growht.pdf')
#plt.close()
def plot_stat(x,y,title,filename,cut=False):
    if cut:
        y = map(lambda i: i/1000,y)
        title += ' (x1000)'
    fig = plt.figure()
    plt.clf()
    ax = plt.axes(FIG_AXES2)
    plt.plot_date(x,y,fmt='k.-')
    plt.ylabel(title)
    plt.grid(True)
    labels = ax.get_xticklabels()
    plt.setp(labels, rotation=30)
    ax.xaxis.set_major_locator(major_loc)
    ax.xaxis.set_minor_locator(minor_loc)
    ax.xaxis.set_major_formatter(fmtr)
    plt.savefig(filename)
    plt.close()

plot_stat(days,edges,'Links','gowalla_growth_links.pdf',cut=True)
plot_stat(days,nodes,'Nodes','gowalla_growth_nodes.pdf',cut=True)
#plot_stat(days,degrees,'Average node degree','gowalla_growth_degree.pdf')


plt.figure()
plt.clf()
ax = plt.axes(FIG_AXES2)
log_nodes = map(math.log,nodes)
log_edges = map(math.log,edges)
a,b = plt.polyfit(log_nodes,log_edges,1)
a2, b2 = plt.polyfit(nodes, edges, 1)
print a,b
print a2, b2
fit = map(lambda i: math.exp(b)*(i**a),nodes)
fit2 = map(lambda i: a2*i + b2, nodes)
print fit
print fit2
plt.loglog(nodes,edges,'ko',mfc='none')
fx = plt.loglog(nodes,fit,'k-',linewidth=1)
fx2 = plt.loglog(nodes,fit2,'k--',linewidth=1)
plt.legend([fx[0], fx2[0]],[
  '$K_t \propto N_t^{%.2f}$'%a,
  '$K_t \propto %.2f N_t$'%a2],
        loc = 'upper left')

loc,lbl = plt.xticks()
lbl = map(lambda i: '%d'%(i/1000), loc)
plt.xticks(loc,lbl)

loc,lbl = plt.yticks()
lbl = map(lambda i: '%d'%(i/1000), loc)
plt.yticks(loc,lbl)

plt.xlabel('$N_t$ (x1000)')
plt.ylabel('$K_t$ (x1000)')
plt.grid(True)
plt.savefig('densification.pdf')
